import json
import os
import random

import tqdm

from ai.utils import utils_file


def data_convert_to_json(accent_scp_path, wav_scp_path, train_path='./output/accent_train.json',
                         dev_path='./output/accent_dev.json'):
    os.makedirs(os.path.dirname(train_path), exist_ok=True)
    dev_ratio = 0.1
    res_train = dict()
    res_dev = dict()
    os.makedirs(os.path.dirname(train_path), exist_ok=True)
    os.makedirs(os.path.dirname(dev_path), exist_ok=True)
    with open(accent_scp_path, 'r', encoding='utf-8') as f:
        accent = f.readlines()
    with open(wav_scp_path, 'r', encoding='utf-8') as f:
        wav = f.readlines()
    assert len(accent) == len(wav)
    num = len(accent)
    res = []
    for index, (wav_str, accent_str) in enumerate(zip(wav, accent)):
        wav_id1, wav_path = wav_str.strip().split()
        wav_id2, wav_accent = accent_str.strip().split()
        assert wav_id1 == wav_id2
        i = random.randint(1, num)
        if i < num * dev_ratio:
            res_dev[wav_id1] = {
                'spk_id': wav_accent,
                'wav': wav_path
            }
        else:
            res_train[wav_id1] = {
                'spk_id': wav_accent,
                'wav': wav_path
            }
    with open(train_path, 'w', encoding='utf-8') as f:
        json.dump(res_train, f, ensure_ascii=False, indent=4)
    with open(dev_path, 'w', encoding='utf-8') as f:
        json.dump(res_dev, f, ensure_ascii=False, indent=4)


def get_little_devset_from_json(json_file_path='./output/accent_dev.json', num=100):
    dict_res = utils_file.load_json_file(json_file_path)
    dict_list = list(dict_res.values())
    random.shuffle(dict_list)
    res_list = dict_list[:num]
    output_path = json_file_path.replace('.json', f'_little_{num}.jsonl')
    utils_file.write_dictlist_to_jsonl(res_list, output_path)


def copy_data_to_my_home(jsonl_file: str):
    """"""
    dic_list = utils_file.load_jsonl_file(jsonl_file)
    all_file_list = [value['wav'] for value in dic_list]
    json_file_name = jsonl_file.strip().split('/')[-1].split('.')[0]
    output_dir = os.path.join('./output/', json_file_name)
    utils_file.makedir_for_file_or_dir(output_dir)
    command = "cp " + " ".join(all_file_list) + f" {output_dir}"
    status = os.system(command)
    # 判断退出状态是否为0，如果为0表示成功，否则表示失败
    if status == 0:
        print(f"successful")
    else:
        print(f"failed")


def data_convert_to_json_for_new_dataset(wav_scp_file: str, label_scp_file: str, output_json_file: str):
    """"""
    prefix = "/home/work_nfs6/qjshao/qjshao_workspace/DATASET/KeSpeech/"
    os.makedirs(os.path.dirname(output_json_file), exist_ok=True)
    res = {}
    wav_dic = utils_file.load_dic_from_scp(wav_scp_file)
    label_dic = utils_file.load_dic_from_scp(label_scp_file)
    for k, v in tqdm.tqdm( label_dic.items(), total=len(label_dic)):
        res[k] = {
            'wav': utils_file.join_path(prefix, wav_dic[k]),
            'spk_id': '普通话' if v == 'Mandarin' else '方言'
        }
    utils_file.write_dict_to_json(res, output_json_file)


if __name__ == "__main__":
    """"""
    data_dir = "/home/work_nfs6/qjshao/qjshao_workspace/DATASET/KeSpeech/Tasks/SubdialectID/test_phase1/"
    data_convert_to_json_for_new_dataset(
        wav_scp_file=utils_file.join_path(data_dir, "wav.scp"),
        label_scp_file=utils_file.join_path(data_dir, "utt2subdialect"),
        output_json_file='./output/accent_test.json')
